import math
import numpy as np

def Lboundary(x, lx):
    res=x
    while x>=lx:
        x-=lx
    while x<0:
        x+=lx
    return x

def boundary(pos, lbox):
    pos[0]=Lboundary(pos[0],lbox[0])
    pos[1]=Lboundary(pos[1],lbox[1])
    pos[2]=Lboundary(pos[2],lbox[2])
    return pos

def anint(x, lx):
    res=x
    if x>0.5*lx:
        res-=lx
    elif x<-0.5*lx:
        res+=lx
    return res

def distancia(lbox,posi,posj):
    dis2=0
    drx=posj[0]-posi[0]
    drx=anint(drx, lbox[0])
    dry=posj[1]-posi[1]
    dry=anint(dry, lbox[1])
    drz=posj[2]-posi[2]
    #drz=anint(drz, lbox[2])
    dis2=drx*drx+dry*dry+drz*drz
    return  [math.sqrt(dis2),drx,dry,drz]

def eKin(vel,diam):
    ek=0.0;
    i=0
    for v in vel:
        ek+=math.pow(diam[i],3.0)*(v[0]*v[0]+v[1]*v[1]+v[2]*v[2])
        i+=1
    return 0.5*ek/len(vel)

def EpotenDip(m,lbox,pos,diam):
    u0=m*m/8.0
    ePot=0.
    n=0
    for i in range(len(pos)):
        for j in range(len(pos)):
            if i!=j:
                [dis, dx,dy,dz]= distancia(lbox,pos[i],pos[j])
                Co= dz/dis
                Co2=Co*Co
                dis3=dis*dis*dis
                potdm = u0*math.pow(diam[i],3.0)*math.pow(diam[j],3.0)*(1.-3.*Co2)/dis3
                if n==0:
                    print 'dis: ',dis, potdm,dz,Co
                ePot+=potdm
                n+=1
    return 0.5*ePot/len(pos)

def EpotenKlin(m,lbox,datos):
    u0=m*m
    kappa=100.0
    ePot=0.
    for i in range(len(datos)):
        for j in range(len(datos)):
            if i!=j:
                [dis, dx,dy,dz]= distancia(lbox,pos[i],pos[j])
                potdm = 3.*u0/kappa*math.exp(kappa*(1-dis))
                    #print p
                    #print i,j,dis*2.0, dis2,potdm,[dis2, dx,dy,dz]
                ePot+=potdm
    return 0.5*ePot/len(datos)

def stressTensor(m,lbox,pos,vel, diam):
    vol=lbox[0]*lbox[1]*lbox[2]
    npart = len(pos)
    sumv=np.zeros((3,3))
    sumf=np.zeros((3,3))
    sigma=np.zeros((3,3))
    for i in range(len(vel)):
        for i1 in range(3):
            for i2 in range(3):
                sumv[i1][i2]+=math.pow(diam[i], 3.0)*vel[i][i1]*vel[i][i2]
    for i in range(npart-1):
        for j in range(i+1,npart):
            force=forceDipij(m,lbox, pos, diam, i,j)
            [dis,dx,dy,dz] = distancia(lbox,pos[i],pos[j])
            dis=[dx,dy,dz]
            for i1 in range(3):
                for i2 in range(3):
                    sumf[i1][i2]+=force[i2]*dis[i1]
                    #if i1==0 and i2==1:
    print 'aaaaaaaaaaaaaaaaaaaaaaaaa',sumf[0][1],vol

    for i1 in range(3):
            for i2 in range(3):
                sigma[i1][i2]=(sumv[i1][i2]+sumf[i1][i2])/vol
    return [sumv,sumf,sigma]

def forceDipij(m,lbox, pos, diam, i,j):
    u0=m*m*math.pow(diam[i],3.0)*math.pow(diam[j],3.0)/8.0
    [dis,dx,dy,dz] = distancia(lbox,pos[i],pos[j])
    Co= dz/dis
    Co2=Co*Co
    fx=3.*u0*(5.*Co2-1.)*dx/math.pow(dis,5.0)
    fy=3.*u0*(5.*Co2-1.)*dy/math.pow(dis,5.0)
    fz=3.*u0*(5.*Co2-3.)*dz/math.pow(dis,5.0)
    return [fx,fy,fz]

def load(name):
    nline=0
    pos=[]
    vel=[]
    force=[]
    diam=[]
    for line in file(name):
        tx = line.replace('\n','')
        val = tx.split(' ')
        res=[]
        for v in val:
            if v!='':
                res.append(v)
        if nline==0:
            t=float(res[1])
            npart=int(res[5])
            lbox=[float(res[7]),float(res[8]),float(res[9])]
        else:
            if(nline<=npart):
                xxx = boundary([float(res[0]),float(res[1]),float(res[2])], lbox)
                pos.append(xxx)
                vel.append([float(res[3]),float(res[4]),float(res[5])])
                force.append([float(res[6]),float(res[7]),float(res[8])])
                diam.append(float(res[9]))
        nline+=1
    return [pos,vel,force, diam,lbox]

u0=100.0
m=math.sqrt(u0)
name = '/home/jonk/eclipse/test/configuracion.dat'

[pos,vel,force, diam,lbox]=load(name)
print len(pos), lbox
ek= eKin(vel,diam)
ep=EpotenDip(m,lbox,pos,diam)
[sumv,sumf,sigma]=stressTensor(m,lbox,pos,vel, diam)
print ek,ep
print sigma[0][1],sigma[0][2],sigma[1][2],sigma[2][2]